package spark_sql;

import bean.User;
import java.util.Arrays;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.expressions.MutableAggregationBuffer;
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructType;

/**
 * @author shihb
 * @date 2020/1/10 22:10
 * 用户自定义聚合函数
 */
public class UdafDemo {

  public static void main(String[] args) {
    //local模式,创建SparkConf对象设定spark的部署环境
    SparkConf sparkConf = new SparkConf().setMaster("local[*]").setAppName("");
    //创建spark上下文对象
    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
    //创建SparkSQL的环境对象
    SparkSession spark = new SparkSession(jsc.sc());


    //创建Rdd
    JavaRDD<User> rdd = jsc.parallelize(Arrays.asList(
        new User(1, "zhangsan", 20),
        new User(2, "lisi", 30),
        new User(3, "wangwu", 40)
    ),1);

    Encoder<User> userEncoder = Encoders.bean(User.class);
    Dataset<User> dataset = spark.createDataset(rdd.rdd(), userEncoder);
    dataset.createOrReplaceTempView("user");

    //创建udaf函数
    MyAvgAgeFunction udaf = new MyAvgAgeFunction();
    //注册udaf聚合函数
    spark.udf().register("avgAge", udaf);



    //使用udaf函数
    Dataset<Row> result = spark.sql("select avgAge(age) from user");
    result.show();
    //释放资源
    spark.stop();
  }

}


/**
 * 用户自定义聚合函数
 * 求age平均值
 */
class MyAvgAgeFunction extends UserDefinedAggregateFunction{

  /**
   * 定义输入的数据类型
   * @return StructType类型
   */
  @Override
  public StructType inputSchema() {
    return new StructType().add("age", DataTypes.LongType);
  }
  /**
   * 计算中的数据类型
   * @return StructType类型
   */
  @Override
  public StructType bufferSchema() {
    return new StructType().add("sum", DataTypes.LongType).add("count", DataTypes.LongType);
  }
  /**
   * 返回值的数据类型
   * @return DataType类型
   */
  @Override
  public DataType dataType() {
    return DataTypes.DoubleType;
  }

  /**
   * 是否稳定
   * @return
   */
  @Override
  public boolean deterministic() {
    return true;
  }


  /**
   * 计算前，缓冲的初始化
   * @param buffer 缓存区
   */
  @Override
  public void initialize(MutableAggregationBuffer buffer) {
    buffer.update(0,0L);
    buffer.update(1,0L);

  }

  /**
   * 新增输入操作
   * @param buffer 缓存区
   * @param row 输入
   */
  @Override
  public void update(MutableAggregationBuffer buffer, Row row) {
    buffer.update(0,buffer.getLong(0)+row.getLong(0));
    buffer.update(1,buffer.getLong(1)+1);
  }

  /**
   * 多个节点的缓冲区合并
   * @param buffer 当前节点缓存区
   * @param other 其他节点缓存区
   */
  @Override
  public void merge(MutableAggregationBuffer buffer, Row other) {
    buffer.update(0,buffer.getLong(0)+other.getLong(0));
    buffer.update(1,buffer.getLong(1)+other.getLong(1));
  }

  /**
   * 计算并返回结果
   * @param buffer 缓存区
   * @return
   */
  @Override
  public Object evaluate(Row buffer) {
    return (buffer.getLong(0)+0.0)/buffer.getLong(1);
  }
}